﻿#include "draw.h"
#include "shared.h"
#include "shape.h"
#include "shader.h"
#include "target.h"
#include "texture.h"
#include "context.h"
#include "font.h"
#include <stack>

// ANIP_UNDONE: remove manual AA & use real AA in ui's main loop
// maybe also wrap up a better multisample-buffer

using namespace anip;

struct Draw::Impl
{
	static Shape
		* square,
		* polyline;

	static Shader
		* lineShader,
		* polylineShader,
		* singleColorShader, 
		* roundedShader, 
		* texShader,
		* hdrTexShader,
		* textShader, 
		* ellipShader,
		* arcShader,
		* shadowRectShader;

	static bool 
		initialized;

	static void init();
};

Shape* Draw::Impl::square = nullptr;
Shape* Draw::Impl::polyline = nullptr;

Shader* Draw::Impl::lineShader = nullptr;
Shader* Draw::Impl::polylineShader = nullptr;
Shader* Draw::Impl::singleColorShader = nullptr;
Shader* Draw::Impl::roundedShader = nullptr;
Shader* Draw::Impl::ellipShader = nullptr;
Shader* Draw::Impl::arcShader = nullptr;
Shader* Draw::Impl::texShader = nullptr;
Shader* Draw::Impl::hdrTexShader = nullptr;
Shader* Draw::Impl::textShader = nullptr;
Shader* Draw::Impl::shadowRectShader = nullptr;

bool Draw::Impl::initialized;

void Draw::Impl::init()
{
	float pts[8] = {
		0.0f, 0.0f,
		1.0f, 0.0f,
		0.0f, 1.0f,
		1.0f, 1.0f
	};

	static const char* vs =
		"#version 410\n"
		"out vec2 pos;"
		"layout(location = 0) in vec2 vPos;"
		"uniform mat3 transform;"
		"uniform vec2 translate;"
		"uniform vec2 size;"
		"void main()"
		"{"
			"gl_Position = vec4(transform * vec3(vPos * size + translate, 1.0f), 1.0f);"
			"pos = vPos * size;"
		"}";

	static const char* linevs =
		"#version 410\n"
		"out float len;"
		"layout(location = 0) in vec2 vPos;"
		"uniform vec2 start;"
		"uniform vec2 dir;"
		"uniform vec2 norm;"
		"uniform mat3 transform;"
		"uniform float dstart;"
		"void main()"
		"{"
			"vec2 mpos = start + vPos.x * dir + (vPos.y - 0.5) * norm;"
			"gl_Position = vec4(transform * vec3(mpos, 1.0f), 1.0f);"
			"len = vPos.x * length(dir) + dstart;"
		"}";

	static const char* polylinevs =
		"#version 410\n"
		"out float len;"
		"layout(location = 0) in vec2 vPos;"
		"layout(location = 1) in float l;"
		"uniform float wid;"
		"uniform mat3 transform;"
		"void main()"
		"{"
			"len = l;"
			"gl_Position = vec4(transform * vec3(vPos, 1.0f), 1.0f);"
		"}";

	static const char* fsLineCol =
		"#version 410\n"
		"in float len;"
		"out vec4 fragColor;"
		"uniform vec4 color;"
		"uniform vec2 dash;"
		"void main()"
		"{"
			"if (mod(len, dash.x) > dash.y) fragColor = vec4(0.0, 0.0, 0.0, 0.0);"
			"else fragColor = color;"
		"}";

	static const char* fsCol =
		"#version 410\n"
		"in vec2 pos;"
		"out vec4 fragColor;"
		"uniform vec4 color;"
		"void main()"
		"{"
			"fragColor = color;"
		"}";

	static const char* fsRounded =
		"#version 410\n"
		"uniform float roundRadius;"
		"uniform vec4 color;"
		"uniform vec2 size;"
		"in vec2 pos;"
		"out vec4 fragColor;"
		"void main()"
		"{"
			"vec2 np = vec2(min(pos.x, size.x - pos.x), min(pos.y, size.y - pos.y));"
			"float a = length(np - vec2(roundRadius)) - roundRadius;"
			"if (np.x < roundRadius && np.y < roundRadius && a > 0.0f)"
				"fragColor = vec4(color.r, color.g, color.b, color.a * (1-a));"
			"else fragColor = color;"
		"}";

	static const char* fsEllip =
		"#version 410\n"
		"in vec2 pos;"
		"out vec4 fragColor;"
		"uniform vec2 size;"
		"uniform vec4 color;"
		"void main()"
		"{"
			"float l = length(abs(vec2(0.5) - pos / size));"
			"if (l < 0.5f)"
				"fragColor = color;"
			"else discard;"
		"}";

	static const char* fsArc =
		"#version 410\n"
		"in vec2 pos;"
		"out vec4 fragColor;"
		"uniform vec2 size;"
		"uniform vec4 color;"
		"uniform vec2 begin_end;"
		"uniform float wid;"
		"void main()"
		"{"
			"vec2 p = pos / size;"
			"float l = length(abs(vec2(0.5) - p));"
			"float m = length(abs(vec2(0.5) - (pos - vec2(wid)) / (size - vec2(2.0 * wid))));"
			"if (l < 0.5 && m > 0.5)"
			"{"
				"float angle = atan(p.y - 0.5, p.x - 0.5);"
				"if (angle < begin_end.x || angle > begin_end.y) discard;"
				"fragColor = color;"
			"}"
			"else discard;"
		"}";

	static const char* fsTex =
		"#version 410\n"
		"in vec2 pos;"
		"out vec4 fragColor;"
		"uniform sampler2D tex;"
		"uniform vec2 size;"
		"uniform vec4 tint;"
		"void main()"
		"{"
			"fragColor = texture(tex, pos / size) * tint;"
		"}";

	static const char* fsHDRTex = R"(
#version 410
#define EPSILON 0.001
in vec2 pos;
out vec4 fragColor;
uniform sampler2D tex;
uniform vec2 size;
uniform vec4 tint;
uniform vec2 markup;
void main()
{
	vec4 col = texture(tex, pos / size) * tint;
	if (markup.x > 0)
	{
		if (col.r < -EPSILON || col.g < -EPSILON || col.b < -EPSILON)
			col *= mod(pos.x+pos.y, markup.x) > (markup.y+EPSILON) ? 1.0 : 0.0;
		if (col.r > 1+EPSILON || col.g > 1+EPSILON || col.b > 1+EPSILON)
			col *= mod(pos.x-pos.y, markup.x) > (markup.y+EPSILON) ? 1.0 : 0.0;
	}
	fragColor = col;
}
)";

	static const char* fsText =
		"#version 410\n"
		"in vec2 pos;"
		"out vec4 fragColor;"
		"uniform sampler2D tex;"
		"uniform vec2 size;"
		"uniform vec4 color;"
		"void main()"
		"{"
			"float r = texture(tex, pos / size).r;"
			"fragColor = color * vec4(1, 1, 1, r);"
		"}";

	static const char* fsShadow =
		"#version 410\n"
		"in vec2 pos;"
		"out vec4 fragColor;"
		"uniform vec2 size;"
		"uniform vec4 color;"
		"uniform float blur;"
		"void main()"
		"{"
			"vec2 np = vec2(min(pos.x, size.x - pos.x), min(pos.y, size.y - pos.y));"
			"float fade = 1;"
			"if (np.x < blur && np.y < blur)"
				"fade = max(1.0 - distance(np, vec2(blur, blur)) / blur, 0);"
			"else if (np.x < blur)"
				"fade = np.x / blur;"
			"else if (np.y < blur)"
				"fade = np.y / blur;"
			"fragColor = color * vec4(1, 1, 1, fade * fade);"
		"}";

	square = new Shape(Shape::Type::TriangleStrip, pts, 4);

	polyline = new Shape(Shape::Type::TriangleStrip);
	polyline->addDataSlot(2); polyline->bindData(0, 0);
	polyline->addDataSlot(1); polyline->bindData(1, 1);

	auto vert = new ShaderPart(ShaderPart::Type::Vertex, vs);
	auto lineVert = new ShaderPart(ShaderPart::Type::Vertex, linevs);
	auto polylineVert = new ShaderPart(ShaderPart::Type::Vertex, polylinevs);

	auto lineColor = new ShaderPart(ShaderPart::Type::Fragment, fsLineCol);
	auto singleColor = new ShaderPart(ShaderPart::Type::Fragment, fsCol);
	auto rounded = new ShaderPart(ShaderPart::Type::Fragment, fsRounded);
	auto texFrag = new ShaderPart(ShaderPart::Type::Fragment, fsTex);
	auto hdrTexFrag = new ShaderPart(ShaderPart::Type::Fragment, fsHDRTex);
	auto ellip = new ShaderPart(ShaderPart::Type::Fragment, fsEllip);
	auto arc = new ShaderPart(ShaderPart::Type::Fragment, fsArc);
	auto textFrag = new ShaderPart(ShaderPart::Type::Fragment, fsText);
	auto shadow = new ShaderPart(ShaderPart::Type::Fragment, fsShadow);

	lineShader = new Shader{ lineVert, lineColor };
	polylineShader = new Shader{ polylineVert, lineColor };
	singleColorShader = new Shader{ vert, singleColor };
	roundedShader = new Shader{ vert, rounded };
	texShader = new Shader{ vert, texFrag };
	hdrTexShader = new Shader{ vert, hdrTexFrag };
	ellipShader = new Shader{ vert, ellip };
	arcShader = new Shader{ vert, arc };
	textShader = new Shader{ vert, textFrag };
	shadowRectShader = new Shader{ vert, shadow };

	Shared::add({ square, polyline, vert, singleColor, rounded, texFrag, ellip, 
		textFrag, shadow, singleColorShader, roundedShader, texShader, 
		hdrTexShader, ellipShader, textShader, shadowRectShader });

	initialized = true;
}

void Draw::unitSquareShape()
{
	//assert(Impl::initialized);
	Impl::square->draw();
}

void anip::Draw::_init()
{
	if (!Impl::initialized)
		Impl::init();
}

void anip::Draw::rect(invec2 pos, invec2 size, const color& fill)
{
	if (fill.a == 0) return;
	//assert(Impl::initialized);
	Impl::singleColorShader->use();
	Impl::singleColorShader->set("translate", pos);
	Impl::singleColorShader->set("transform", 
		Context::currentTarget().transformMatrix());
	Impl::singleColorShader->set("size", size);
	Impl::singleColorShader->set("color", fill);
	Impl::square->draw();
}

void anip::Draw::roundedRect(invec2 pos, invec2 size, float rad, const color& fill)
{
	if (fill.a == 0) return;
	//assert(Impl::initialized);
	Impl::roundedShader->use();
	Impl::roundedShader->set("translate", pos);
	Impl::roundedShader->set("transform",
		Context::currentTarget().transformMatrix());
	Impl::roundedShader->set("size", size);
	Impl::roundedShader->set("color", fill);
	Impl::roundedShader->set("roundRadius", rad);
	Impl::square->draw();
}

void anip::Draw::shadowRect(invec2 pos, invec2 size, const color& fill, float blur)
{
	constexpr float expendBlur = 1.5; // ANIP_FIXME: ?? review

	if (fill.a == 0) return;
	//assert(Impl::initialized);
	Impl::shadowRectShader->use();
	Impl::shadowRectShader->set("translate", pos - vec2{ blur, blur });
	Impl::shadowRectShader->set("transform", 
		Context::currentTarget().transformMatrix());
	Impl::shadowRectShader->set("size", size + vec2{ blur * 2, blur * 2 });
	Impl::shadowRectShader->set("color", fill);
	Impl::shadowRectShader->set("blur", blur * expendBlur);
	Impl::square->draw();
}

void anip::Draw::ellipse(invec2 pos, invec2 size, const color& fill)
{
	if (fill.a == 0) return;
	//assert(Impl::initialized);
	Impl::ellipShader->use();
	Impl::ellipShader->set("translate", pos);
	Impl::ellipShader->set("transform", 
		Context::currentTarget().transformMatrix());
	Impl::ellipShader->set("size", size);
	Impl::ellipShader->set("color", fill);
	Impl::square->draw();
}

void anip::Draw::rectOutline(invec2 pos, invec2 size, float wid, const color& col)
{
	if (col.a == 0 || wid == 0) return;
	// or use polyline?
	Draw::line(pos, pos + size * vec2{ 1, 0 }, wid, col);
	Draw::line(pos, pos + size * vec2{ 0, 1 }, wid, col);
	Draw::line(pos + size * vec2{ 1, 0 }, pos + size * vec2{ 1, 1 }, wid, col);
	Draw::line(pos + size * vec2{ 0, 1 }, pos + size * vec2{ 1, 1 }, wid, col);
}

void anip::Draw::roundedRectOutline(invec2 pos, invec2 size, float rad, float wid, const color& col)
{
	if (col.a == 0 || wid == 0) return;
	if (rad == 0) 
	{ 
		rectOutline(pos, size, wid, col); 
		return;
	}

	Draw::line(pos + vec2{ rad, 0 }, pos + vec2{ size.x - rad, 0 }, wid, col);
	Draw::line(pos + vec2{ 0, rad }, pos + vec2{ 0, size.y - rad }, wid, col);
	Draw::line(pos + vec2{ rad, size.y }, pos + vec2{ size.x - rad, size.y }, wid, col);
	Draw::line(pos + vec2{ size.x, rad }, pos + vec2{ size.x, size.y - rad }, wid, col);

	vec2 rsize{ rad * 2, rad * 2 };
	Draw::arc(pos, rsize, -3.14f, -1.57f, wid, col);
	Draw::arc(pos + vec2{ size.x - rad * 2, 0 }, rsize, -1.57f, 0, wid, col);
	Draw::arc(pos + vec2{ 0, size.y - rad * 2 }, rsize, 1.57f, 3.14f, wid, col);
	Draw::arc(pos + size - rsize, rsize, 0, 1.57f, wid, col);
}

void anip::Draw::ellipseOutline(invec2 pos, invec2 size, float wid, const color& col)
{
	if (col.a == 0 || wid == 0) return;
	//assert(Impl::initialized);
	Impl::arcShader->use();
	Impl::arcShader->set("translate", pos);
	Impl::arcShader->set("transform",
		Context::currentTarget().transformMatrix());
	Impl::arcShader->set("size", size);
	Impl::arcShader->set("begin_end", vec2{ -4, 4 });
	Impl::arcShader->set("wid", wid);
	Impl::arcShader->set("color", col);
	Impl::square->draw();
}

void anip::Draw::arc(invec2 pos, invec2 size, 
	float start, float end, float wid, const color& col)
{
	if (col.a == 0 || wid == 0) return;
	//assert(Impl::initialized);
	Impl::arcShader->use();
	Impl::arcShader->set("translate", pos);
	Impl::arcShader->set("transform",
		Context::currentTarget().transformMatrix());
	Impl::arcShader->set("size", size);
	Impl::arcShader->set("begin_end", vec2{ start, end });
	Impl::arcShader->set("wid", wid);
	Impl::arcShader->set("color", col);
	Impl::square->draw();
}

void anip::Draw::line(invec2 p1, invec2 p2, float wid, 
	const color& col, const DashOption& dash)
{
	if (col.a == 0 || wid == 0) return;
	//assert(Impl::initialized);
	Impl::lineShader->use();
	Impl::lineShader->set("start", p1);
	Impl::lineShader->set("dir", p2 - p1);
	Impl::lineShader->set("norm", math::normal(p2 - p1) * wid);
	if (dash.solidLength > 0)
	{
		Impl::lineShader->set("dash", vec2{ dash.unitLength, dash.solidLength });
		Impl::lineShader->set("dstart", dash.start);
	}
	else
		Impl::lineShader->set("dash", vec2{ 1.0f, 1.0f });
	Impl::lineShader->set("transform", 
		Context::currentTarget().transformMatrix());
	Impl::lineShader->set("color", col);
	Impl::square->draw();
}

void anip::Draw::polyline(const vec2* pts, int n, float wid,
	const color& col, const DashOption& dash)
{
	if (pts == nullptr || n < 1 || col.a == 0 || wid == 0) return;
	//assert(Impl::initialized);

	// ANIP_FIXME: this algorithm is flawed, go implement
	// https://blog.scottlogic.com/2019/11/18/drawing-lines-with-webgl.html
	// in the future
	auto ps = new vec2[n * 2];
	auto len = new float[n * 2];
	vec2 norm = math::normal(pts[1] - pts[0]);
	wid -= 1;
	ps[0] = pts[0] + norm * wid * 0.5f;
	ps[1] = pts[0] - norm * wid * 0.5f;
	len[0] = len[1] = dash.start;
	for (int i = 1; i < n; i++)
	{
		norm = math::normal(pts[i + 1] - pts[i]);
		ps[i*2] = pts[i] + norm * wid * 0.5f;
		ps[i*2+1] = pts[i] - norm * wid * 0.5f;
		len[i*2] = len[i*2+1] = len[i*2-1] + math::distance(pts[i], pts[i-1]);
	}

	Impl::polyline->getData(0).set(n * 2, (float*)ps);
	Impl::polyline->getData(1).set(n * 2, (float*)len);

	delete[] ps;
	delete[] len;

	Impl::polylineShader->use();
	if (dash.solidLength > 0)
		Impl::polylineShader->set("dash", vec2{ dash.unitLength, dash.solidLength });
	else
		Impl::polylineShader->set("dash", vec2{ 1.0f, 1.0f });
	Impl::polylineShader->set("transform", 
		Context::currentTarget().transformMatrix());
	Impl::polylineShader->set("color", col);

	Impl::polyline->draw();
}

void anip::Draw::crossMarker(invec2 pos, float size, float wid, const color& color)
{
	line(pos - vec2{ size, size }, pos + vec2{ size, size }, wid, color);
	line(pos - vec2{ -size, size }, pos + vec2{ -size, size }, wid, color);
}

void anip::Draw::texture(invec2 pos, const Texture2D& tex, const color& tint)
{
	texture(pos, tex.size(), tex, tint);
}

void anip::Draw::texture(invec2 pos, invec2 size,
	const Texture2D& tex, const color& tint)
{
	if (tint.a == 0) return;
	//assert(Impl::initialized);
	Impl::texShader->use();
	Impl::texShader->set("translate", pos);
	Impl::texShader->set("transform", 
		Context::currentTarget().transformMatrix());
	Impl::texShader->set("size", size);
	Impl::texShader->set("tex", tex, 0);
	Impl::texShader->set("tint", tint);
	Impl::square->draw();
}

void anip::Draw::textureHDR(invec2 pos, invec2 size, 
	const Texture2D& tex, const color& tint, bool markOutOfRange)
{
	if (tint.a == 0) return;
	//assert(Impl::initialized);
	Impl::hdrTexShader->use();
	Impl::hdrTexShader->set("translate", pos);
	Impl::hdrTexShader->set("transform",
		Context::currentTarget().transformMatrix());
	Impl::hdrTexShader->set("size", size);
	Impl::hdrTexShader->set("tex", tex, 0);
	Impl::hdrTexShader->set("tint", tint);
	if (markOutOfRange)
		Impl::hdrTexShader->set("markup", 6.0f, 0.5f);
	else
		Impl::hdrTexShader->set("markup", -1.0f, -1.0f);
	Impl::square->draw();
}

void anip::Draw::text(invec2 pos, invec2 size,
	const Texture2D& tex, const color& col)
{
	// ANIP_UNDONE: snapping to pixel is needed?

	if (col.a == 0) return;
	//assert(Impl::initialized);
	Impl::textShader->use();
	Impl::textShader->set("translate", pos);
	Impl::textShader->set("transform", 
		Context::currentTarget().transformMatrix());
	Impl::textShader->set("size", size);
	Impl::textShader->set("tex", tex, 0);
	Impl::textShader->set("color", col);
	Impl::square->draw();
}

void anip::Draw::text(invec2 pos, const wstring& str,
	const color& col, const Font& f, invec2 align)
{
	if (col.a == 0) return;
	text(pos, str, col, f, align, TextRenderer::Options());
}

void anip::Draw::text(invec2 pos, const wstring& str,
	const color& col, const Font& f, invec2 align, const TextRenderer::Options& opt)
{
	assert(!f.isNull());
	auto layout = TextRenderer::layoutPixelPerfectString(str, f, opt);
	vec2 size = layout->drawingSize();
	Texture2D* tex = TextRenderer::renderString(*layout);
	delete layout;
	if (tex == nullptr) return;
	text(pos - align * size, size, *tex, col);
	tex->finalize();
	delete tex;
}